﻿#include <vector>
#include <iostream>
#include <map>
using namespace std;

class Solution_subarraysWithKDistinct {
public:
    int subarraysWithKDistinct(vector<int>& A, int K) {
        if (K==0 || A.size() == 0)
            return 0;
		const int len = A.size();
        // 統計最多K的所有子數組
        map<int, int> mapKMost;
        // 統計最多(k-1)的所有子數組
        map<int, int> mapLessK;
        int leftKMost = 0;
        int leftLessK = 0;
        int right = 0;
        int ans = 0;
        while (right < len) {
            // add to both map
            const int numR = A[right];
			mapKMost[numR] = (mapKMost.count(numR) > 0 ? mapKMost[numR] : 0) + 1;
			mapLessK[numR] = (mapLessK.count(numR) > 0 ? mapLessK[numR] : 0) + 1;
            // deal if size is above the limit k
            while (mapKMost.size() > K) {
                const int num = A[leftKMost];
                const int cnt = mapKMost[num];
                if (cnt <= 1)
                    mapKMost.erase(num);
                else
                    mapKMost[num] = cnt - 1;
                leftKMost++;
            }
            // deal if size is above the limit k-1
            while (mapLessK.size() > K - 1) {
                const int num = A[leftLessK];
                const int cnt = mapLessK[num];
                if (cnt <= 1)
                    mapLessK.erase(num);
                else
                    mapLessK[num] = cnt - 1;
                leftLessK++;
            }
            // (right-leftKMost+1)+(right-leftLessK+1) => leftLessK - leftKMost
            ans += leftLessK - leftKMost;
            // right pointer add
            right++;
        }
        return ans;
    }
};

int main()
{
   Solution_subarraysWithKDistinct so;
   vector<int> num = { 1, 2, 1, 2, 3 };
	cout << "want:7, ans = " << so.subarraysWithKDistinct(num, 2) << endl;

   vector<int> num1 = { 1, 2, 1, 3, 4 };
	cout << "want:3, ans = " << so.subarraysWithKDistinct(num1, 3) << endl;

   vector<int> num2 = { 2, 2, 3, 3, 3, 1, 3 };
   cout << "want:13, ans = " << so.subarraysWithKDistinct(num2, 2) << endl;

   vector<int> num3 = { 2, 1, 1, 1, 2 };
   cout << "want:8, ans = " << so.subarraysWithKDistinct(num3, 1) << endl;

   return 0;
}
